import os
import sys
import pandas as pd
import auto_zdock
import ChianChange
import rosetta_relax
index = pd.read_csv(
	"./input/index.csv",
	sep=","
)
chain_list = ["A","B","C","D","E","F","G","H","I","J","K","L","M","N"]
receptor_from_csv = list(index["receptor"])
ligand_from_csv = list(index["ligand"])
receptor_antibody_or_not_from_csv = list(index["receptor_antibody"])

receptor = []
ligand = []
receptor_antibody_or_not = []

for i in receptor_from_csv:
	if type(i) == type("str"):
		receptor.append(i)
for i in ligand_from_csv:
	if type(i) == type("str"):
		ligand.append(i)
for i in receptor_antibody_or_not_from_csv:
	if type(i) == type("str"):
		receptor_antibody_or_not.append(i)

if (len(receptor) != len(ligand)) | (len(receptor) != len(receptor_antibody_or_not)):
    print("输入的’index.csv‘里面 receptor、ligand、receptor_antibody 数量不等请检查！")
    sys.exit(0)

if len(receptor) == 0:
	print("输入的’index.csv‘ 内容为空！")
	sys.exit(0)

for i in range(len(receptor)):
	os.system("cp ./input/{}.pdb ./".format(receptor[i]))
	os.system("cp ./input/{}.pdb ./".format(ligand[i]))
	#检查链名是不是重复
	receptor_chains = []
	ligand_chains = []

	rec = open("./{}.pdb".format(receptor[i]),"r")
	rec_lines = [j for j in rec]
	for j in range(len(rec_lines)):
		if rec_lines[j][0:4] == 'ATOM':
			receptor_chains.append(rec_lines[j][21].strip())
	receptor_chains_unique = list(set(receptor_chains))

	lig = open("./{}.pdb".format(ligand[i]),"r")
	lig_lines = [j for j in lig]
	for j in range(len(lig_lines)):
		if lig_lines[j][0:4] == 'ATOM':
			ligand_chains.append(lig_lines[j][21].strip())
	ligand_chains_unique = list(set(ligand_chains))

	# print(receptor_chains_unique)
	# print("".join(receptor_chains_unique))
	# print(ligand_chains_unique)
	# print("".join(ligand_chains_unique))
	len_rec_chain = len(receptor_chains_unique)
	len_lig_chain = len(ligand_chains_unique)
	rec_chain_str = "".join(chain_list[0:len_rec_chain])
	lig_chain_str = "".join(chain_list[len_rec_chain:len_rec_chain+len_lig_chain])
	ChianChange.chainchange(receptor[i],"".join(receptor_chains_unique),rec_chain_str)
	ChianChange.chainchange(ligand[i],"".join(ligand_chains_unique),lig_chain_str)
	# print(rec_chain_str)
	# print(lig_chain_str)

	print("zdock")
	auto_zdock.auto_zdock(receptor[i]+"_chain_changed.pdb",ligand[i]+"_chain_changed.pdb",receptor_antibody_or_not[i])

# rosetta_relax.relax()
# os.system("rm ./*.pdb")






